import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch_geometric.nn import GATConv
from torch_geometric.data import Data
import numpy as np

class ResourceGNN(nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, heads=2):
        super(ResourceGNN, self).__init__()
        self.gat1 = GATConv(in_channels, hidden_channels, heads=heads)
        self.gat2 = GATConv(hidden_channels * heads, out_channels, heads=1)

    def forward(self, x, edge_index):
        x = F.elu(self.gat1(x, edge_index))
        x = self.gat2(x, edge_index)
        return x

class UserBehaviorTransformer(nn.Module):
    def __init__(self, input_dim, num_heads, hidden_dim, num_layers):
        super(UserBehaviorTransformer, self).__init__()
        self.encoder_layer = nn.TransformerEncoderLayer(
            d_model=input_dim,
            nhead=num_heads,
            dim_feedforward=hidden_dim,
            activation='relu'
        )
        self.transformer = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)
        self.fc = nn.Linear(input_dim, input_dim)

    def forward(self, x):
        x = self.transformer(x)
        x = self.fc(x[:, -1, :])  # 取最后一个时间步
        return x

class RecommendationSystem(nn.Module):
    def __init__(self, gnn, transformer):
        super(RecommendationSystem, self).__init__()
        self.gnn = gnn
        self.transformer = transformer
        self.similarity = nn.CosineSimilarity(dim=-1)

    def forward(self, user_history, item_graph):
        # GNN 计算资源表征
        item_embeddings = self.gnn(item_graph.x, item_graph.edge_index)

        # Transformer 计算用户行为表征
        user_embedding = self.transformer(user_history)

        # 计算用户与所有资源的相似度
        scores = self.similarity(user_embedding.unsqueeze(0), item_embeddings)
        return scores

if __name__ == "__main__":
    # 生成示例数据 (学习资源图)
    num_resources = 10
    resource_features = torch.rand((num_resources, 16))  # 16 维特征
    edge_index = torch.randint(0, num_resources, (2, 20))  # 随机生成 20 条边
    item_graph = Data(x=resource_features, edge_index=edge_index)

    # 生成用户行为序列 (用户最近浏览的 5 个资源)
    user_history = torch.rand((5, 16))  # 5 个时间步，每个时间步 16 维特征

    # 初始化模型
    gnn = ResourceGNN(in_channels=16, hidden_channels=32, out_channels=16)
    transformer = UserBehaviorTransformer(input_dim=16, num_heads=2, hidden_dim=32, num_layers=2)
    recommender = RecommendationSystem(gnn, transformer)

    # 计算推荐分数
    scores = recommender(user_history.unsqueeze(0), item_graph)
    top_k = torch.argsort(scores, descending=True)[:5]  # 取前 5 个推荐资源
    print("推荐的学习资源索引:", top_k.tolist())
